package com.cscec5b.inspection.config;

import com.cscec5b.inspection.entity.ClientApp;
import com.cscec5b.inspection.repository.ClientRepository;
import com.nimbusds.jose.crypto.RSASSAVerifier;
import com.nimbusds.jwt.SignedJWT;
import jakarta.servlet.FilterChain;
import jakarta.servlet.ServletException;
import jakarta.servlet.http.HttpServletRequest;
import jakarta.servlet.http.HttpServletResponse;
import org.springframework.http.HttpHeaders;
import org.springframework.stereotype.Component;
import org.springframework.util.StreamUtils;
import org.springframework.web.filter.OncePerRequestFilter;

import java.io.IOException;
import java.nio.charset.StandardCharsets;
import java.security.MessageDigest;
import java.security.interfaces.RSAPublicKey;
import java.text.ParseException;
import java.time.Instant;
import java.util.Base64;
import java.util.Optional;

import lombok.extern.slf4j.Slf4j;

@Slf4j
@Component
public class JwtAuthFilter extends OncePerRequestFilter {
    private final KeyManager keyManager;
    private final AppProperties props;
    private final ClientRepository clientRepo;

    public JwtAuthFilter(KeyManager keyManager, AppProperties props, ClientRepository clientRepo) {
        this.keyManager = keyManager; this.props = props; this.clientRepo = clientRepo;
    }

    @Override
    protected boolean shouldNotFilter(HttpServletRequest request) {
        String path = request.getRequestURI();
        if ("OPTIONS".equalsIgnoreCase(request.getMethod())) return true;

        org.springframework.util.AntPathMatcher m = new org.springframework.util.AntPathMatcher();

        // 0) 若命中“必须鉴权的前缀”，直接要求鉴权（覆盖白名单）
        for (String p : props.getProtectedPrefixes()) {
            if (m.match(p, path)) return false; // 不放行
        }

        // 1) 其余按白名单匹配
        for (String p : props.getAuthExcludePaths()) {
            if (m.match(p, path)) return true;
        }
        return false;
    }

    private static String fp(RSAPublicKey k) {
        try {
            var md = MessageDigest.getInstance("SHA-256");
            return Base64.getEncoder().encodeToString(md.digest(k.getEncoded())).substring(0, 16);
        } catch (Exception e) { return "n/a"; }
    }

    @Override
    protected void doFilterInternal(HttpServletRequest req, HttpServletResponse resp, FilterChain chain)
            throws ServletException, IOException {
        String auth = req.getHeader(HttpHeaders.AUTHORIZATION);
        if (auth == null || !auth.startsWith("Bearer ")) { resp.sendError(401, "missing bearer token"); return; }
        // === 1) 解析 JWT 并按 kid 选公钥验签 ===
        SignedJWT jwt;
        try {
            jwt = SignedJWT.parse(auth.substring(7));
            String kid = jwt.getHeader().getKeyID();

            // 支持 kid（KeyManager 里我们加过 publicKey(String kid) 重载，先忽略 kid 返回同一把公钥也行）
            RSAPublicKey pub = (kid != null) ? keyManager.publicKey(kid) : keyManager.publicKey();
            if (pub == null) { log.warn("JWT reject: kid={} no pubkey", kid); resp.sendError(401, "kid not found"); return; }

            boolean ok = jwt.verify(new RSASSAVerifier(pub));
            if (!ok) { log.warn("JWT reject: signature invalid, kid={}, pub.fp={}", kid, fp(pub)); resp.sendError(401, "invalid jwt"); return; }

            if (Instant.now().isAfter(jwt.getJWTClaimsSet().getExpirationTime().toInstant())) {
                log.warn("JWT reject: expired");
                resp.sendError(401, "jwt expired"); return;
            }
        } catch (Exception e) {
            log.warn("JWT reject: parse/verify error", e);
            resp.sendError(401, "bad jwt"); return;
        }

        // === 2) 头部基础校验（timestamp 兼容毫秒/秒）===
        String ts = req.getHeader("X-Timestamp");
        String rid = req.getHeader("X-Request-Id");
        if (ts == null || rid == null) { resp.sendError(400, "missing X-Timestamp or X-Request-Id"); return; }

        long t;
        try { t = Long.parseLong(ts); } catch (NumberFormatException e) { resp.sendError(400, "bad timestamp"); return; }
        // 如果是毫秒，转换为秒
        long tsSec = (t > 1_000_000_000_000L) ? (t / 1000) : t;
        long nowSec = Instant.now().getEpochSecond();
        if (Math.abs(nowSec - tsSec) > props.getMaxClockSkewSeconds()) {
            resp.sendError(403, "timestamp skew too large"); return;
        }

        // 可选：请求体签名（HMAC）
        if (props.isRequireSignature()) {
            String sign = req.getHeader("X-Signature");
            if (sign == null) { resp.sendError(403, "missing X-Signature"); return; }
            String clientId;
            try { clientId = jwt.getJWTClaimsSet().getSubject(); }
            catch (ParseException e) { resp.sendError(401, "no sub"); return; }

            Optional<ClientApp> cli = clientRepo.findByClientId(clientId);
            if (cli.isEmpty()) { resp.sendError(401, "unknown client"); return; }

            String body = new String(StreamUtils.copyToByteArray(req.getInputStream()), StandardCharsets.UTF_8);
            String base = ts + "\n" + body;
            String expect = SignatureUtil.hmacSha256Hex(cli.get().getHmacSecret(), base);
            if (!expect.equalsIgnoreCase(sign)) { resp.sendError(403, "bad signature"); return; }

            req = new CachedBodyRequest(req, body); // 让后续还能读 body
        }

        try { req.setAttribute("clientId", jwt.getJWTClaimsSet().getSubject()); }
        catch (ParseException e) { resp.sendError(401, "bad claims"); return; }

        chain.doFilter(req, resp);
    }

    // ========= 辅助类 =========
    static class SignatureUtil {
        static String hmacSha256Hex(String secret, String data) {
            try {
                javax.crypto.Mac mac = javax.crypto.Mac.getInstance("HmacSHA256");
                mac.init(new javax.crypto.spec.SecretKeySpec(secret.getBytes(StandardCharsets.UTF_8), "HmacSHA256"));
                byte[] raw = mac.doFinal(data.getBytes(StandardCharsets.UTF_8));
                StringBuilder sb = new StringBuilder(raw.length * 2);
                for (byte b : raw) sb.append(String.format("%02x", b));
                return sb.toString();
            } catch (Exception e) { throw new RuntimeException(e); }
        }
    }
    static class CachedBodyRequest extends jakarta.servlet.http.HttpServletRequestWrapper {
        private final byte[] cached;
        CachedBodyRequest(HttpServletRequest r, String body) {
            super(r); this.cached = body.getBytes(StandardCharsets.UTF_8);
        }
        @Override public jakarta.servlet.ServletInputStream getInputStream() {
            return new jakarta.servlet.ServletInputStream() {
                private int idx = 0;
                @Override public int read() { return (idx < cached.length) ? (cached[idx++] & 0xff) : -1; }
                @Override public boolean isFinished() { return idx >= cached.length; }
                @Override public boolean isReady() { return true; }
                @Override public void setReadListener(jakarta.servlet.ReadListener readListener) {}
            };
        }
    }
}
